import numpy as np
from sklearn.model_selection import KFold, GroupKFold, StratifiedKFold, train_test_split

__all__ = ['init_folds', 'train_test_split']


def init_folds(metadata, img_group_id_colname=None, img_class_colname=None):
    """
    Initialzies the cross-validation splits.

    Parameters
    ----------
    img_group_id_colname : str or None
        Column in `metadata` that is used to create cross-validation splits.
        If not None, then images that have the same group_id are never in train and validation.
    img_class_colname : str or None
        Column in `metadata` that is used to create cross-validation splits. If not none,
        splits are stratifed to ensure the same distribution of `img_class_colname` in train and validation.

    Returns
    -------

    """
    n_folds = 2
    fold = -1
    skip_train = 1
    if img_group_id_colname is not None:
        gkf = GroupKFold(n_folds)
        if img_class_colname is not None:
            class_col_name = getattr(metadata, img_class_colname, None)
        else:
            class_col_name = None
        splitter = gkf.split(X=metadata,
                             y=class_col_name,
                             groups=getattr(metadata, img_group_id_colname))
    else:
        if img_class_colname is not None:
            skf = StratifiedKFold(n_folds)
            splitter = skf.split(X=metadata,
                                 y=getattr(metadata, img_class_colname, None))
        else:
            kf = KFold(n_folds)
            splitter = kf.split(X=metadata)

    cv_split = []
    for fold_id, (train_ind, val_ind) in enumerate(splitter):

        if fold != -1 and fold_id != fold:
            continue

        np.random.shuffle(train_ind)
        train_ind = train_ind[::skip_train]

        cv_split.append((fold_id,
                         metadata.iloc[train_ind],
                         metadata.iloc[val_ind]))
    return cv_split
